package com.example.demo.graph;

import com.alibaba.cloud.ai.graph.CompiledGraph;
import com.alibaba.cloud.ai.graph.KeyStrategy;
import com.alibaba.cloud.ai.graph.StateGraph;
import com.alibaba.cloud.ai.graph.action.AsyncNodeAction;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor;

import org.springframework.ai.chat.prompt.ChatOptions;


import org.springframework.context.annotation.Bean;
import org.springframework.stereotype.Component;

import java.util.HashMap;
import java.util.Map;

import static com.alibaba.cloud.ai.graph.StateGraph.END;
import static com.alibaba.cloud.ai.graph.StateGraph.START;
import com.alibaba.cloud.ai.graph.node.LlmNode;
import org.springframework.ai.chat.messages.AssistantMessage;

@Component
public class GraphBuilder {

    private final ChatClient.Builder chatClientBuilder;

    public GraphBuilder(ChatClient.Builder chatClientBuilder) {
        this.chatClientBuilder = chatClientBuilder;
    }

    @Bean
    public CompiledGraph buildGraph() throws Exception {
        //ChatClient chatClient = ChatClient.builder(chatModel).defaultAdvisors(new SimpleLoggerAdvisor()).build();

        ChatClient gpt4o = chatClientBuilder
                .defaultOptions(ChatOptions.builder()
                        .model("gpt-4o")
                        .temperature(0.7)
                        .build())
                .defaultAdvisors(new SimpleLoggerAdvisor())
                .build();

        /*ChatClient o3Client = chatClientBuilder
                .defaultOptions(ChatOptions.builder()
                        .model("o3")
                        .temperature(0.7)
                        .build())
                .defaultAdvisors(new SimpleLoggerAdvisor())
                .build();*/

        // new stateGraph
        // todo: Add some non-node variable names, such as sys and conversation variables. Format as sys_xxx.
        StateGraph stateGraph = new StateGraph(() -> {
            Map<String, KeyStrategy> strategies = new HashMap<>();
            strategies.put("sys_query", (o1, o2) -> o2);
            strategies.put("LLMNode3_text", (o1, o2) -> o2);
            strategies.put("end1_output", (o1, o2) -> o2);
            strategies.put("LLMNode1_text", (o1, o2) -> o2);
            strategies.put("startNode1_input", (o1, o2) -> o2);
            strategies.put("LLMNode2_text", (o1, o2) -> o2);
            return strategies;
        });
        // add nodes
        // EndNode [ 1757503742203 ]
        stateGraph.addNode("end1", AsyncNodeAction
                .node_async(state -> Map.of("end1_output", Map.of("result", state.value("LLMNode3_text").orElse("")))));// —— LlmNode [1760414088279] ——

        stateGraph.addNode("LLMNode1", AsyncNodeAction.node_async(state -> {
            String input1 = state.value("startNode1_input").orElse("").toString();
            LlmNode LLMNode1 = LlmNode.builder()
                    .systemPromptTemplate(
                            "你是资深的B2B商业分析师，请根据【公司画像】生成潜在客户画像，该画像将用于后续AI通过关键词去寻找真实目标公司，因此内容必须具体、可操作，便于关键词提炼。\n一、核心任务与输入依据​\n1，【公司画像】：{startNode1_input}\n2，核心任务：根据上述已有信息，推理该公司的潜在客户画像，禁止使用无关信息或臆造内容。​\n\n二、输出要求​\n每个维度的描述不超过500字，要求中文，详细且具象，便于后续AI提取“行业名”“企业规模”“痛点词”等关键词。\n输出语气分析型，注重逻辑推理与商业可行性。\n输出结构：必须严格符合下列 JSON 结构，缺失信息时填空字符串 \"\"；\n{\n\"companyOverview\": 【维度一：潜在客户公司概况】\n\"请基于该公司服务对象的典型特征，从以下角度推理并整合成一段分析性文字：\n 1.**所属行业及细分领域**：明确行业大类（如制造业）及具体子类（如汽车零部件二级供应商、食品饮料OEM代工厂），避免泛称“工业企业”。\n2. **企业类型与性质**：是否为上市公司、集团子公司、家族企业、外资工厂、跨境电商卖家等。\n3. **规模特征**：根据该公司产品定价、部署复杂度、客户案例等，推断目标客户的典型员工规模（如100–500人）、年营收区间（如5000万–3亿）、分支机构数量等。\n4. **地域与市场覆盖**：是否聚焦特定区域（如长三角制造业集群、东南亚出口型企业）、是否服务海外市场（如欧美合规要求高的客户）。\n5. **采购决策特征**：是IT部门主导（如SaaS工具）、生产部门驱动（如工业设备）、还是高管战略采购（如数字化转型项目）。\n示例逻辑：'因该公司提供中等价位的MES系统（单价20–50万元），且案例多为华东地区电子组装厂，可推断其客户多为员工200–800人的中型制造企业，具备一定自动化基础但缺乏数据整合能力，痛点集中在生产追溯与OEE提升。'\",\n\"productService\":【维度二：潜在客户的产品】\n\"请基于【公司画像】的产品与服务部分，推导其合作客户的主营产品，禁止推导出服务“\n输出示例：”1)产品大类1（产品细类1、产品细类2...）,2)产品大类2（产品细类1、产品细类2...）...“\n}")
                    .userPromptTemplate("以下是【公司画像】：\n{startNode1_input}\n\n" +
                            "请基于上述信息，生成符合要求的潜在客户画像JSON。")
                    .params(Map.of("startNode1_input", input1))
                    .chatClient(gpt4o)
                    .outputKey("LLMNode1_text").build();

            Map<String, Object> result = LLMNode1.apply(state);
            String key = "LLMNode1_text";
            Object object = result.get(key);
            if (object instanceof AssistantMessage && ((AssistantMessage) object).getText() != null) {
                return Map.of(key, ((AssistantMessage) object).getText());
            }
            return Map.of(key, object != null ? object.toString() : "unknown");
        }));

        // —— LlmNode [1760414153291] ——
        stateGraph.addNode("LLMNode2", AsyncNodeAction.node_async(state -> {
            String input2 = state.value("LLMNode1_text").orElse("").toString();
            LlmNode LLMNode2 = LlmNode.builder()
                    .systemPromptTemplate(
                            "角色定位：\n你是一名熟悉 Google 搜索语义匹配逻辑的 B2B 市场分析师，擅长根据潜客画像提取真实可用于 Google 搜索的高价值关键词。\n\n" +
                                    "【任务目标】\n" +
                                    "从【潜客画像】的维度二产品部分中，提炼出适用于 Google 搜索目标公司/品牌的产品。\n" +
                                    "这些产品要能帮助 AI 在 Google 上找到最符合潜客特征的公司官网或企业介绍页面。\n\n" +
                                    "【输出要求】\n" +
                                    "- 输出五个以内的关键词；\n" +
                                    "- 每个关键词必须是从【潜客画像】的产品中提取的产品；\n" +
                                    "- 只能输出产品相关词；\n\n" +
                                    "【输出结构】\n" +
                                    "仅输出以下 JSON 结构：\n" +
                                    "{\n\"keywords\": [\"产品1\", \"产品2\", ...]\n}\n\n" +
                                    "特殊说明：若输入的潜客画像报告完全无有效信息，输出时“keywords”字段填空字符串 \"\"；\n" +
                                    "涉及“报告”“文档”等表述时，无需替换，仅需确保关键词为英文。")
                    .userPromptTemplate("{LLMNode1_text}")
                    .params(Map.of("LLMNode1_text", input2))
                    .chatClient(gpt4o)
                    .outputKey("LLMNode2_text")
                    .build();

            // 3️⃣ 执行推理
            Map<String, Object> result = LLMNode2.apply(state);

            // 4️⃣ 提取结果
            String key = "LLMNode2_text";
            Object object = result.get(key);
            String output = (object instanceof AssistantMessage)
                    ? ((AssistantMessage) object).getText()
                    : (object != null ? object.toString() : "unknown");

            return Map.of(key, output.trim());
        }));


        stateGraph.addNode("LLMNode3", AsyncNodeAction.node_async(state -> {
            String input3 = state.value("LLMNode1_text").orElse("").toString();
            String input4 = state.value("LLMNode2_text").orElse("").toString();
            LlmNode LLMNode3 = LlmNode.builder()
                    .systemPromptTemplate(
                            "你是一个专业的信息抽取与结构化处理助手。\n" +
                                    "你需要接收两部分内容，并从中精准提取信息，然后严格按照指定的 JSON 格式输出完整内容：\n\n" +
                                    "【输入说明】\n1️⃣ 第一部分为潜在客户画像（companyOverview、productService 等字段）\n" +
                                    "2️⃣ 第二部分为从潜客画像中提炼出的关键词信息\n\n" +
                                    "【输出要求】\n1. 必须严格遵守以下 JSON 结构；\n" +
                                    "2. 所有字段必须存在，即使无内容也用 \"\" 填充；\n" +
                                    "3. 不要添加注释、解释或其他多余内容；\n" +
                                    "4. 输出仅限于 JSON 对象，不包含任何前缀或后缀；\n\n" +
                                    "【输出结构】\n{\n" +
                                    "  \"companyOverview\": \"\",\n" +
                                    "  \"productService\": \"\",\n" +
                                    "  \"keywords\": [\"英文关键词1\", \"英文关键词2\", ...]\n" +
                                    "}\n\n" +
                                    "【注意】\n1. 所有字段必须存在，无内容时用空字符串 \"\" 填充；\n" +
                                    "2. 输出必须是合法 JSON，可直接被解析；\n" +
                                    "3. 禁止输出任何解释性文字。")
                    .userPromptTemplate("【潜客画像】:\n{LLMNode1_text}\n\n【关键词提取结果】:\n{LLMNode2_text}")
                    .params(Map.of(
                            "LLMNode1_text", input3,
                            "LLMNode2_text", input4
                    ))
                    .chatClient(gpt4o)
                    .outputKey("LLMNode3_text")
                    .build();

            // 3️⃣ 执行调用
            Map<String, Object> result = LLMNode3.apply(state);

            // 4️⃣ 统一输出格式
            String key = "LLMNode3_text";
            Object object = result.get(key);
            String output = (object instanceof AssistantMessage)
                    ? ((AssistantMessage) object).getText()
                    : (object != null ? object.toString() : "unknown");

            return Map.of(key, output.trim());
        }));


        // add edges
        stateGraph.addEdge("LLMNode3", "end1");
        stateGraph.addEdge("LLMNode2", "LLMNode3");
        stateGraph.addEdge(START, "LLMNode1");
        stateGraph.addEdge("LLMNode1", "LLMNode2");
        stateGraph.addEdge("end1", END);

        return stateGraph.compile();
    }

}
